from numpy import zeros, matrix
from element import element


class BeamElement(Element):
    def __init__(self, element_id, youngs_modulus, moment_of_inertia, length):
        self.element_id = element_id
        self.youngs_modulus = youngs_modulus
        self.moment_of_inertia = moment_of_inertia
        self.length = length

    def calculate_stiffness_matrix(self, E, I, L):
        return (E*I/L**3)*matrix([[12.0, 6.0*L,  -12.0,    6.0*L],
                                  [6.0*L, 4.0*L**2, -6.0*L, 2.0*L**2],
                                  [-12.0,   -6.0*L,   12.0,   -6.0*L],
                                  [6.0*L, 2.0*L**2, -6.0*L,  4.0*L*2]])
        
    def set_stiffness(self):
        self.stiffness_matrix = self.calculate_stiffness_matrix(self.youngs_modulus, 
                                                                self.moment_of_inertia, self.length)
